//
// Created by apple on 2022/9/25.
//

#include "matrix.h"
#include "dot.h"

Matrix::Matrix(int nRow, int nCol) {
    this->_nRow = nRow;
    this->_nCol = nCol;
    data = new float[nRow * nCol];
}

Matrix *Matrix::setRandom() {
    std::default_random_engine engine(clock());
    std::uniform_real_distribution<float> dist(1e-2, 1);
    for (size_t i = 0; i < _nRow * _nCol; ++i) {
        float temp = dist(engine);
        data[i] = temp;
    }
    return this;
}


Matrix *Matrix::setZero() {
    for (size_t i = 0; i < _nRow * _nCol; ++i)
        data[i] = 0;
    return this;
}

void Matrix::print(const std::string &end) const {
    std::cout << "Row:" << _nRow << " | Col:" << _nCol << std::endl;
    for (int i = 0; i < _nRow; ++i) {
        for (int j = 0; j < _nCol; ++j) {
            auto temp = data[i * _nCol + j];
            std::cout.precision(3);
            if (temp >= 0)
                std::cout << " " << temp << "\t";
            else
                std::cout << temp << " ";
        }
        std::cout << end;
    }
}

Matrix *Matrix::transpose() const {
    auto m = new Matrix(_nCol, _nRow);
    for (int i = 0; i < _nRow; ++i)
        for (int j = 0; j < _nCol; ++j)
            m->data[j * _nRow + i] = this->data[i * _nCol + j];
    return m;
}

Matrix *matmulBaseline(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    float *a = A->data, *b = B->data, *c = C->data;
    for (size_t i = 0; i < A->nRow(); ++i) {
        for (size_t j = 0; j < B->nCol(); ++j) {
            float temp = 0;
            for (size_t k = 0; k < A->nCol(); ++k)
                temp += a[i * A->nCol() + k] * b[k * B->nCol() + j];
            c[i * B->nCol() + j] = temp;
        }
    }
    return C;
}


void matmulBaselineThreadTask(Matrix *A, Matrix *B, Matrix *C, int i_start, int i_end) {
    float *a = A->data, *b = B->data, *c = C->data;
    for (size_t i = i_start; i < i_end; ++i) {
        for (size_t j = 0; j < B->nCol(); ++j) {
            float temp = 0;
            for (size_t k = 0; k < A->nCol(); ++k)
                temp += a[i * A->nCol() + k] * b[k * B->nCol() + j];
            c[i * B->nCol() + j] = temp;
        }
    }

}

Matrix *matmulBaselineMultiThread(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    C->setZero();
    int thread_num = std::min(A->nRow(), THREAD_NUM);
    int interval = std::ceil(A->nRow() / thread_num);
    std::thread pool[thread_num];
    for (int i = 0; i < thread_num; ++i)
        pool[i] = std::thread(&matmulBaselineThreadTask, A, B, C, i * interval,
                              std::min((i + 1) * interval, A->nRow()));
    for (int i = 0; i < thread_num; ++i)
        pool[i].join();
    return C;

}

Matrix *matmulIKJ(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    C->setZero();
    float *a = A->data, *b = B->data, *c = C->data;
    for (size_t i = 0; i < A->nRow(); ++i)
        for (size_t k = 0; k < A->nCol(); ++k) {
            float s = a[i * A->nCol() + k];
            for (size_t j = 0; j < B->nCol(); ++j)
                c[i * B->nCol() + j] += s * b[k * B->nCol() + j];
        }
    return C;
}

void matmulIJKThreadTask(Matrix *A, Matrix *B, Matrix *C, int i_start, int i_end) {
    float *a = A->data, *b = B->data, *c = C->data;
    for (size_t i = i_start; i < i_end; ++i)
        for (size_t k = 0; k < A->nCol(); ++k) {
            float s = a[i * A->nCol() + k];
            for (size_t j = 0; j < B->nCol(); ++j)
                c[i * B->nCol() + j] += s * b[k * B->nCol() + j];
        }
}


Matrix *matmulIKJMultiThread(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    C->setZero();
    int thread_num = std::min(A->nRow(), THREAD_NUM);
    int interval = std::ceil(A->nRow() / thread_num);
    std::thread pool[thread_num];
    for (int i = 0; i < thread_num; ++i)
        pool[i] = std::thread(matmulIJKThreadTask, A, B, C, i * interval, std::min((i + 1) * interval, A->nRow()));
    for (int i = 0; i < thread_num; ++i)
        pool[i].join();
    return C;
}


Matrix *matmulBlock(Matrix *A, Matrix *B) {
    // in ikj order
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    C->setZero();
    float *a = A->data, *b = B->data, *c = C->data;
    for (int ii = 0; ii < A->nRow(); ii += BLOCK_SIZE)
        for (int jj = 0; jj < B->nCol(); jj += BLOCK_SIZE)
            for (int kk = 0; kk < A->nCol(); kk += BLOCK_SIZE)
                for (int i = ii; i < std::min(ii + BLOCK_SIZE, A->nRow()); i++)
                    for (int k = kk; k < std::min(kk + BLOCK_SIZE, A->nCol()); k++) {
                        float s = a[i * A->nCol() + k];
                        for (int j = jj; j < std::min(jj + BLOCK_SIZE, B->nCol()); j++)
                            c[i * B->nCol() + j] += s * b[k * B->nCol() + j];
                    }
    return C;
}

void matmulBlockThreadTask(Matrix *A, Matrix *B, Matrix *C, const std::vector<int> &ii_data) {
    // in ikj order
    float *a = A->data, *b = B->data, *c = C->data;
    for (auto ii: ii_data)
        for (int jj = 0; jj < B->nCol(); jj += BLOCK_SIZE)
            for (int kk = 0; kk < A->nCol(); kk += BLOCK_SIZE)
                for (int i = ii; i < std::min(ii + BLOCK_SIZE, A->nRow()); i++)
                    for (int k = kk; k < std::min(kk + BLOCK_SIZE, A->nCol()); k++) {
                        float s = a[i * A->nCol() + k];
                        for (int j = jj; j < std::min(jj + BLOCK_SIZE, B->nCol()); j++)
                            c[i * B->nCol() + j] += s * b[k * B->nCol() + j];
                    }
}


Matrix *matmulBlockMultiTread(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    C->setZero();

    int block_num = std::ceil((double) A->nRow() / BLOCK_SIZE);
    int thread_num = std::min(THREAD_NUM, block_num);
    int block_num_per_thread = std::ceil(block_num / thread_num);
    std::thread pool[thread_num];
    int ii = 0;
    for (int i = 0; i < thread_num; ++i) {
        std::vector<int> ii_data;
        for (; ii < std::min((i + 1) * block_num_per_thread, block_num); ++ii)
            ii_data.push_back(ii * BLOCK_SIZE);

        pool[i] = std::thread(&matmulBlockThreadTask, A, B, C, ii_data);
    }
    for (size_t i = 0; i < thread_num; ++i)
        pool[i].join();
    return C;
}

Matrix *matmulWinergrad(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    auto B_T = B->transpose();
    float *a = A->data, *b = B_T->data, *c = C->data;
    for (size_t i = 0; i < A->nRow(); ++i) {
        for (size_t j = 0; j < B_T->nRow(); ++j) {
            auto ap = a + i * A->nCol();
            auto bp = b + j * B_T->nCol();
            c[i * B->nCol() + j] = dotWinergrad(ap, bp, A->nCol());
        }
    }
    delete B_T;
    return C;
}

Matrix *matmulIKJWinergrad(Matrix *A, Matrix *B) {
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    C->setZero();
    float *a = A->data, *b = B->data, *c = C->data;
    int k_len = A->nCol();
    if (k_len % 2 == 1) {
        k_len -= 1;
        for (size_t i = 0; i < A->nRow(); ++i)
            for (size_t j = 0; j < B->nCol(); ++j)
                c[i * B->nCol() + j] += a[(i + 1) * A->nCol() - 1] * b[(A->nCol() - 1) * B->nCol() + j];
    }

    for (size_t i = 0; i < A->nRow(); ++i) {
        for (size_t k = 0; k < k_len; k += 2) {
            float a_odd = a[i * A->nCol() + k];
            float a_even = a[i * A->nCol() + k + 1];
            for (size_t j = 0; j < B->nCol(); ++j) {
                auto b_odd = b[k * B->nCol() + j];
                auto b_even = b[(k + 1) * B->nCol() + j];
                c[i * B->nCol() + j] += (a_odd + b_even) * (a_even + b_odd) - a_odd * a_even - b_odd * b_even;
            }
        }
    }
    return C;
}

Matrix *matmulBlockWinergrad(Matrix *A, Matrix *B) {
    // in ikj order
    assert(A->nCol() == B->nRow());
    auto C = new Matrix(A->nRow(), B->nCol());
    float *a = A->data, *b = B->data, *c = C->data;
    C->setZero();

    int k_len = A->nCol();
    if (k_len % 2 == 1) {
        --k_len;
        for (size_t i = 0; i < A->nRow(); ++i)
            for (size_t j = 0; j < B->nCol(); ++j)
                c[i * B->nCol() + j] += a[(i + 1) * A->nCol() - 1] * b[(A->nCol() - 1) * B->nCol() + j];
    }

    for (int ii = 0; ii < A->nRow(); ii += BLOCK_SIZE)
        for (int jj = 0; jj < B->nCol(); jj += BLOCK_SIZE)
            for (int kk = 0; kk < k_len; kk += BLOCK_SIZE)
                for (int i = ii; i < std::min(ii + BLOCK_SIZE, A->nRow()); i++)
                    for (int k = kk; k < std::min(kk + BLOCK_SIZE, k_len); k+=2) {
                        float a_odd = a[i * A->nCol() + k];
                        float a_even = a[i * A->nCol() + k + 1];
                        for (int j = jj; j < std::min(jj + BLOCK_SIZE, B->nCol()); j++) {
                            auto b_odd = b[k * B->nCol() + j];
                            auto b_even = b[(k + 1) * B->nCol() + j];
                            c[i * B->nCol() + j] +=
                                    (a_odd + b_even) * (a_even + b_odd) - a_odd * a_even - b_odd * b_even;
                        }
                    }
    return C;
}
